import cv2
import numpy as np
import torch


class CenterEncoder:
    """centernet的目标编码器"""

    def __init__(self, input_size, num_cls, stride=4):
        super(CenterEncoder, self).__init__()
        self.input_size = input_size
        self.target_size = [int(x / stride) for x in input_size]
        self.stride = stride
        self.num_cls = num_cls

    def run(self, label):
        # heatmap目标
        t_hm = np.zeros([*self.input_size, self.num_cls], dtype=np.float32)
        # 存在性目标
        t_exist = np.zeros([*self.target_size, self.num_cls], dtype=np.float32)
        # bounding-box 目标
        t_bbox = np.zeros([*self.target_size, 4], dtype=np.float32)
        # 是否忽略
        t_ignore = np.zeros(self.target_size, dtype=np.float32)
        t_bbox[:] = np.nan
        t_ignore[:] = np.nan

        if len(label) > 0:
            xs = (label[:, 0] + label[:, 2]) / 2
            ys = (label[:, 1] + label[:, 3]) / 2
            ws = np.abs(label[:, 2] - label[:, 0])
            hs = np.abs(label[:, 3] - label[:, 1])
            cols = (xs / self.stride).astype(np.int32)
            rows = (ys / self.stride).astype(np.int32)
            types = label[:, 4].astype(np.int32)
            ignored = label[:, 5].astype(np.int8)
            rotated = label[:, 6].astype(np.int8)

            for i in range(len(label)):
                if (
                    rows[i] >= 0
                    and rows[i] < self.target_size[0]
                    and cols[i] >= 0
                    and cols[i] < self.target_size[1]
                ):
                    if ignored[i] > 0:
                        t_ignore[rows[i], cols[i]] = 1.0
                    else:
                        draw_gaussian(
                            t_hm[:, :, types[i]],
                            int(xs[i]),
                            int(ys[i]),
                            gaussian_radius(ws[i], hs[i]),
                        )
                        t_exist[rows[i], cols[i], types[i]] = 1.0
                    if rotated[i] == 0:
                        t_bbox[rows[i], cols[i]] = np.array(
                            [xs[i], ys[i], ws[i], hs[i]], dtype=np.float32
                        )
                else:
                    print(label[i])

        t_hm = cv2.resize(t_hm, dsize=(self.target_size[1], self.target_size[0]))
        if len(t_hm.shape) < 3:
            t_hm = np.expand_dims(t_hm, -1)

        t_hm = np.maximum(t_hm, t_exist)
        targets = np.concatenate((t_bbox, t_hm, np.expand_dims(t_ignore, -1)), axis=-1)
        targets = torch.from_numpy(targets)
        return targets


def draw_gaussian(hm, col, row, w_h_radius, k=1):
    # 构建高斯特征
    w_h_radius = [int(_) for _ in w_h_radius]
    guass_radius = [_ * 2 + 1 for _ in w_h_radius]
    sigmaw, sigmah = [max(1, _ / 3) for _ in w_h_radius]
    gaussian = get_gaussian_2d_kernel(guass_radius, sigmaw, sigmah)
    height, width = hm.shape[0:2]

    left, right = min(col, w_h_radius[0]), min(width - col, w_h_radius[0] + 1)
    top, bottom = min(row, w_h_radius[1]), min(height - row, w_h_radius[1] + 1)

    masked_hm = hm[row - top : row + bottom, col - left : col + right]
    masked_gaussian = gaussian[
        w_h_radius[1] - top : w_h_radius[1] + bottom,
        w_h_radius[0] - left : w_h_radius[0] + right,
    ]
    if min(masked_gaussian.shape) > 0 and min(masked_hm.shape) > 0:
        np.maximum(masked_hm, masked_gaussian * k, out=masked_hm)
    return


def get_gaussian_2d_kernel(shape, sigmaw, sigmah):
    # 构建指定的2d高斯核
    n, m = [(ss - 1.0) / 2.0 for ss in shape]
    n, m = int(n), int(m)
    y, x = np.ogrid[-m : m + 1, -n : n + 1]
    h = np.exp(-(x * x) / (2 * sigmaw * sigmaw) - (y * y) / (2 * sigmah * sigmah))
    h[h < np.finfo(h.dtype).eps * h.max()] = 0
    return h


def gaussian_radius(w, h):
    # 高斯核半径
    w_r = 9 * np.power(np.sqrt(w * h) / 10, 0.8) * np.sqrt(w / (w + h))
    h_r = 9 * np.power(np.sqrt(w * h) / 10, 0.8) * np.sqrt(h / (w + h))
    return w_r, h_r
